# Copyright (c) 2024, Shanghai Iluvatar CoreX Semiconductor Co., Ltd.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
#

from typing import Callable, Union

from .object import iscallable

expr_mapping = dict(
    eq=lambda a, b: a == b,
    ne=lambda a, b: a != b,
    ge=lambda a, b: a >= b,
    le=lambda a, b: a <= b,
    gt=lambda a, b: a > b,
    lt=lambda a, b: a < b,
)


def get_cmp_func(cmp: Union[str, Callable]) -> Callable:
    if iscallable(cmp):
        return cmp

    return expr_mapping[cmp]


def compare_with_op(cmp_op, a, b):
    cmp_op = get_cmp_func(cmp_op)
    return cmp_op(a, b)
